import traceback

from LLM import *
from api_input_tools import get_api_input
from tools import *
from API import *
from utils import match_court_representative


def check_parentheses(string):
    if '(' in string and ')' in string:
        return True
    else:
        return False


def check_court_name(string):
    if '省' in string and "市" in string:
        return True
    else:
        return False


def extract_court_representative(question, data=None):
    PROMPT = f"""
    获取问题中的案号。没有则输出“None”。
    问题："{question}"
    直接输出json，可以被Python json.loads函数解析。不作解释，不作答：
    ```json
    {{
        "案号": ""
    }}
    ```
    """
    response = LLM(PROMPT)
    case_number = prase_json_from_response(response)['案号']
    if case_number is not None:
        court_representative = match_court_representative(case_number)
        return court_representative, court_representative


def retrieve(question, data=None):
    api_name, api_args = get_tools_response(question)

    try:
        if 'case_num' in api_args:
            api_args['case_num'] = api_args['case_num'].replace('（', '(').replace('）', ')')
            case_num = api_args['case_num']
            # 有些问题的案号格式不正确，比如少了(或者)，这里需要判断之后利用LLM修正案号的格式
            if not check_parentheses(case_num):
                api_args['case_num'] = completion_case_num(case_num)
            api_args = get_api_input(api_name, api_args['case_num'])
        elif 'court_name' in api_args:
            court_name = api_args['court_name']
            if not check_court_name(court_name):
                api_args['court_name'] = completion_court_name(court_name)
            api_args = get_api_input(api_name, api_args['court_name'])
        elif 'court_representative' in api_args:
            api_args = get_api_input(api_name, api_args['court_representative'])
        # 有时候需要调用get_lawfirm_log接口，但是tools会返回get_lowfirm_info接口，所以这里虽然都是获取律所名称，但是也区分开来
        elif 'lawfirm_name' in api_args:
            api_args = get_api_input(api_name, api_args['lawfirm_name'])
        elif 'lawfirm_name_log' in api_args:
            api_args = get_api_input(api_name, api_args['lawfirm_name_log'])
        elif 'address' in api_args:
            api_args = get_api_input(api_name, api_args['address'])
        elif 'province_city_area' in api_args:
            # 先将一个完整的地址拆分为省份、城市、区县
            province_city_area = get_province_city_area(api_args['province_city_area'])
            api_args = get_api_input(api_name, province_city_area["省份"], province_city_area["城市"],
                                     province_city_area["区县"])
        elif 'company_name' in api_args:
            api_args = get_api_input(api_name, api_args['company_name'])
        elif 'credit_code' in api_args:
            api_args = get_api_input(api_name, api_args['credit_code'])
        elif 'company_name' in api_args:
            api_args = get_api_input(api_name, api_args['company_name'])
    except:
        pass
    try:
        ori_answer = API(api_name=api_name, args=api_args)
    except Exception as e:
        # 获取异常的堆栈跟踪信息
        traceback_info = traceback.format_exc()
        # 打印堆栈跟踪信息
        print(f"请求API发生异常：{traceback_info}")
    # answer = refine_answer(question, ori_answer)
    return ori_answer, ori_answer


# add by heyong:增加“次”关键字，否则大模型在总结提炼是无法识别到次数
def stat(question, data=None):
    return f"据统计，共{len(data)}次", None


def order(question, data=None):
    PROMPT = f"""
分析所给问题，给出排序属性与最终数量。
问题：{question}
----
请按照以下json格式进行输出，可以被Python json.loads函数解析。不回答问题，不作任何解释，不输出其他任何信息。
对于查询类问题，补充原有属性的具体属性名称。
```json
{{
    "排序属性": "",
    "最终数量": ""
}}
```
"""
    import heapq
    def top_k_elements_with_indices(lst, k):
        indexed_lst = [(val, idx) for idx, val in enumerate(lst)]
        top_k = heapq.nlargest(k, indexed_lst)
        top_k_values, top_k_indices = zip(*top_k)
        return list(top_k_indices), list(top_k_values)

    response = LLM(PROMPT)
    feature = prase_json_from_response(response)

    name_list = []
    value_list = []
    data = [data] if not isinstance(data, list) else data
    for res in data:
        api_name, api_args = get_tools_response(list(res.values())[0] + '的' + feature['排序属性'])
        api_response = API(api_name=api_name, args=api_args)
        name_list.append(api_response)
        value = api_response[feature['排序属性']]
        if value is None:
            value = 0
        value_list.append(float(value))

    idx, _ = top_k_elements_with_indices(value_list, int(feature['最终数量']))
    target = [name_list[i] for i in idx]
    ori_answer = f"排序后的结果：{target}"
    answer = refine_answer(question, ori_answer)
    return answer, ori_answer


def summary(question, data=None):
    return refine_answer(question, data), None


def multi_retrieve(question, data=None):
    PROMPT = f"""
分析所给问题，给出查询属性，中文。
问题：{question}
数据表属性1：标题, 案号, 文书类型, 原告, 被告, 原告律师, 被告律师, 案由, 审理法条依据, 涉案金额, 判决结果, 胜诉方, 文件名
数据表属性2：关联上市公司股票代码, 关联上市公司股票简称, 关联上市公司全称, 上市公司关系, 上市公司参股比例, 上市公司投资金额, 公司名称
----
例子：这些子公司的上市公司参股比例分别是多少？
{{
    "查询属性": "上市公司参股比例"
}}
----
请按照以下json格式进行输出，可以被Python json.loads函数解析。不回答问题，不作任何解释，不输出其他任何信息。
```json
{{
    "查询属性": ""
}}
```
"""
    response = LLM(PROMPT)
    feature = prase_json_from_response(response)
    answer_list = []
    data = [data] if not isinstance(data, list) else data
    for res in data:
        # api_name, api_args = get_tools_response(res + '\n' + question.split('\n')[-1])
        # api_name, api_args = get_tools_response(list(res.values())[0] + '\n' + question.split('\n')[-1])
        api_name, api_args = get_tools_response(list(res.values())[0])

        try:
            api_args['case_num'] = api_args['case_num'].replace('（', '(').replace('）', ')')
        except:
            pass
        api_response = API(api_name=api_name, args=api_args)
        answer_list.append(f"{list(res.values())[0]}的{feature['查询属性']}为{api_response[feature['查询属性']]}")
    return str(answer_list), answer_list


# 暂且写成hard code，自行改进
def filter_list(question, data=None):
    PROMPT = f"""
分析所给问题，给出问题分类，并提取公司名称。
问题：{question}
----
1: 如“华仁药业股份有限公司控股的子公司，超过50%的有几家？”
2: 如“大众交通（集团）股份有限公司中，投资超5000万的子公司有多少家？”
3: 如“熊猫乳品集团股份有限公司的全资控股子公司有几家？”
----
请按照以下json格式进行输出，可以被Python json.loads函数解析。不回答问题，不作任何解释，不输出其他任何信息。
对于查询类问题，补充原有属性的具体属性名称。
```json
{{
    "问题类别序号": "",
    "公司名称": ""
}}
```
"""
    response = LLM(PROMPT)
    feature = prase_json_from_response(response)
    data = [data] if not isinstance(data, list) else data
    sub_com = []
    category = int(feature["问题类别序号"]) if isinstance(feature["问题类别序号"], str) else feature["问题类别序号"]
    if category == 1:
        for res in data:
            sub = API(api_name='get_sub_company_info', args={'company_name': res['公司名称']})
            if sub['上市公司参股比例'] is not None and float(sub['上市公司参股比例']) >= 50.0:
                sub_com.append(sub)
    elif category == 2:
        for res in data:
            sub = API(api_name='get_sub_company_info', args={'company_name': res['公司名称']})
            if sub['上市公司投资金额'] is not None:
                cash = sub['上市公司投资金额']
                if '万' in cash:
                    cash = float(cash[:-1])
                elif '亿' in cash:
                    cash = float(cash[:-1]) * 10000
                else:
                    cash = float(cash)
                if cash >= 5000:
                    sub_com.append(sub)
    else:
        for res in data:
            sub = API(api_name='get_sub_company_info', args={'company_name': res['公司名称']})
            if sub['上市公司参股比例'] is not None and float(sub['上市公司参股比例']) == 100:
                sub_com.append(res['公司名称'])

    answer = f"满足条件的子公司数量为{len(sub_com)}"
    return answer, sub_com


# 暂且写成hard code，自行改进
def calculate_cash(question, data=None):
    data = [data] if not isinstance(data, list) else data
    cash = 0
    for res in data:
        sub = API(api_name='get_sub_company_info', args={'company_name': res['公司名称']})
        if sub['上市公司投资金额'] is not None:
            sub_cash = sub['上市公司投资金额']
            if '万' in sub_cash:
                sub_cash = float(sub_cash[:-1]) / 10000
            elif '亿' in sub_cash:
                sub_cash = float(sub_cash[:-1])
            else:
                sub_cash = float(sub_cash[:-1]) / 10000
            cash += sub_cash
    return f"共投资{cash}亿元", cash
